// Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::{collections::HashSet, io::Write, ops::RangeInclusive, path::{Path, PathBuf}};

use bytes::TryGetError;
use log::{error, warn};
use tokio::fs::create_dir_all;
use walkdir::WalkDir;

use crate::{deploy::MergeMode, Arch};

use super::{BufExt, BufMutExt};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DeployStat {
    pub mode: Option<MergeMode>,
    pub arch: Arch,
    pub pkgs: Vec<String>,
}

pub struct DeployStats {
    pub stats: Vec<DeployStat>,
}

struct FindDeployStat {
    read_dir: tokio::fs::ReadDir,
}

#[derive(Debug, thiserror::Error)]
pub enum DeployStatsError {
    #[error("failed to read deploy stat directory '{0}'")]
    DirReadError(PathBuf, #[source] std::io::Error),
    #[error("failed to create deploy stat directory '{0}'")]
    DirCreateError(#[source] std::io::Error),
    #[error("failed to read deploy stat file '{0}'")]
    FileReadError(PathBuf, #[source] std::io::Error),
    #[error("failed to write deploy stat file '{0}'")]
    FileWriteError(PathBuf, #[source] std::io::Error),
    #[error("failed to update deploy stat symlink '{0}' to '{1}'")]
    UpdateSymlinkError(PathBuf, String, #[source] std::io::Error),
    #[error("current deploy stat id not exists")]
    CurrentIdNotExists,
    #[error("invalid merge mode '{0}'")]
    InvalidMergeMode(u8),
    #[error("invalid arch '{0}'")]
    InvalidArch(u8),
    #[error("invalid deploy stat data")]
    InvalidData(#[from] TryGetError),
    #[error("failed to find deploy stat file")]
    FindDeployStatError(#[source] std::io::Error),
    #[error("failed to remove deploy stat file '{0}'")]
    RemoveFileError(PathBuf, #[source] std::io::Error),
}

#[derive(Debug, thiserror::Error)]
pub enum ParseFileNameError {
    #[error("invalid deploy stat id")]
    InvalidId,
    #[error("invalid deploy stat extension")]
    InvalidExtension,
}

impl DeployStat {
    fn write<B: BufMutExt>(&self, mut buf: B) {
        if let Some(mode) = self.mode {
            buf.put_bool(true);
            buf.put_merge_mode(mode);
        } else {
            buf.put_bool(false);
        }
        buf.put_arch(self.arch);
        buf.put_strs_with_len(self.pkgs.iter().map(|s| s.as_str()), self.pkgs.len() as u32);
    }

    pub fn read<B: BufExt>(mut buf: B) -> Result<Self, DeployStatsError> {
        let undeploy = !buf.try_get_bool()?;
        let mode = if undeploy {
            None
        } else {
            Some(buf.try_get_u8()?
                .try_into()
                .map_err(DeployStatsError::InvalidMergeMode)?)
        };
        let arch = buf.try_get_u8()?
            .try_into()
            .map_err(DeployStatsError::InvalidArch)?;
        let pkgs = buf.try_get_strings_with_len()?;
        Ok(DeployStat { mode, arch, pkgs })
    }

    pub fn dump<W: Write>(&self, mut w: W) -> Result<(), std::io::Error> {
        writeln!(w, "mode: {}", self.mode.map(|m| m.as_str()).unwrap_or("undeploy"))?;
        writeln!(w, "arch: {}", self.arch)?;
        writeln!(w, "pkgs:")?;
        for pkg in self.pkgs.iter() {
            writeln!(w, "  - {}", pkg)?;
        }
        Ok(())
    }

    fn id_to_file_name(id: usize) -> String {
        format!("{}.stat", id)
    }

    fn id_try_from_file_name(file_name: &str) -> Result<usize, ParseFileNameError> {
        if let Some(file_name) = file_name.strip_suffix(".stat") {
            file_name.parse().map_err(|_| ParseFileNameError::InvalidId)
        } else {
            Err(ParseFileNameError::InvalidExtension)
        }
    }

    async fn save<P: AsRef<Path>>(&self, dir: P, id: usize) -> Result<(), DeployStatsError> {
        let mut buf = Vec::new();
        self.write(&mut buf);

        let path = dir.as_ref().join(Self::id_to_file_name(id));
        tokio::fs::write(&path, buf).await
            .map_err(|e| DeployStatsError::FileWriteError(path.clone(), e))?;
        Ok(())
    }

    pub async fn save_after_current<P: AsRef<Path>>(&self, root: P) -> Result<(), DeployStatsError> {
        let dir = DeployStats::deploy_stat_dir(root);
        create_dir_all(&dir).await.map_err(DeployStatsError::DirCreateError)?;
        let current_stat_path = DeployStats::current_stat_path(&dir);
        let current_id = DeployStats::read_current_id(&current_stat_path).await?.unwrap_or(0);
        self.save(&dir, current_id + 1).await?;
        DeployStats::update_current_id(&dir, current_id + 1).await?;
        DeployStats::remove_after(&dir, current_id + 1).await?;
        Ok(())
    }

    pub async fn load_id_range<P: AsRef<Path>>(dir: P, range: RangeInclusive<usize>) -> Result<Vec<Self>, DeployStatsError> {
        let mut stats = vec![];
        for file_name in range.map(Self::id_to_file_name) {
            let file_path = dir.as_ref().join(file_name);
            let content = tokio::fs::read_to_string(&file_path).await
                .map_err(|e| DeployStatsError::FileReadError(file_path.clone(), e))?;
            let stat = DeployStat::read(content.as_bytes())?;
            stats.push(stat);
        }

        Ok(stats)
    }

}

impl DeployStats {
    pub fn deploy_stat_root() -> PathBuf {
        dirs::data_local_dir().unwrap_or(PathBuf::from("~/.local/share"))
            .join("ferrium/deploy")
    }

    pub fn find_deploy_dirs() -> impl Iterator<Item = PathBuf> {
        WalkDir::new(Self::deploy_stat_root())
            .into_iter()
            .filter_map(|e| e.ok())
            .filter(|e| e.file_type().is_symlink() && e.file_name() == "current.stat")
            .filter_map(|e| e.path().parent().map(|p| p.to_path_buf()))
    }

    pub fn deploy_stat_dir<P: AsRef<Path>>(root: P) -> PathBuf {
        let root = root.as_ref().strip_prefix("/").unwrap_or(root.as_ref());
        Self::deploy_stat_root().join(root)
    }

    fn current_stat_path<P: AsRef<Path>>(dir: P) -> PathBuf {
        dir.as_ref().join("current.stat")
    }

    pub async fn collect_deployed_pkg_specs() -> Result<HashSet<String>, DeployStatsError>  {
        let mut deployed_pkgs = HashSet::new();
        for dir in DeployStats::find_deploy_dirs() {
            if let Some((deploy_stats, _)) = DeployStats::load_current(&dir).await? {
                for stat in deploy_stats.stats.iter() {
                    if stat.mode.is_some() {
                        deployed_pkgs.extend(stat.pkgs.iter().cloned());
                    } else {
                        for pkg in stat.pkgs.iter() {
                            deployed_pkgs.remove(pkg);
                        }
                    }
                }
            }
        }
            
        Ok(deployed_pkgs)
    }

    async fn read_current_id<P: AsRef<Path>>(path: P) -> Result<Option<usize>, DeployStatsError> {
        let path = path.as_ref();
        match tokio::fs::read_link(path).await {
            Ok(link) => {
                let file_name = link.as_os_str().to_string_lossy();
                match DeployStat::id_try_from_file_name(file_name.as_ref()) {
                    Ok(id) => Ok(Some(id)),
                    Err(err) => {
                        warn!("invalid deploy stat file '{}': {}", path.display(), err);
                        Ok(None)
                    },
                }
            },
            Err(err) => match err.kind() {
                std::io::ErrorKind::NotFound => Ok(None),
                _ => Err(DeployStatsError::FileReadError(path.to_path_buf(), err)),
            },
        }
    }

    pub async fn update_current_id<P: AsRef<Path>>(dir: P, id: usize) -> Result<(), DeployStatsError> {
        let dir = dir.as_ref();
        let symlink_path = Self::current_stat_path(dir);
        let current_id = Self::read_current_id(&symlink_path).await?.unwrap_or(0);

        if current_id == id {
            return Ok(());
        }
        
        let new_symlink = DeployStat::id_to_file_name(id);
        Self::update_symlink(symlink_path.as_path(), &new_symlink).await
            .map_err(|e| DeployStatsError::UpdateSymlinkError(symlink_path.clone(), new_symlink, e))?;
        Ok(())
    }

    async fn update_symlink(origin: &Path, target: &str) -> Result<(), tokio::io::Error> {
        match tokio::fs::remove_file(origin).await {
            Ok(_) => Ok(()),
            Err(e) if e.kind() == tokio::io::ErrorKind::NotFound => {
                Ok(())
            },
            Err(e) => Err(e),
        }?;
        tokio::fs::symlink(target, origin).await?;
        Ok(())
    }

    async fn remove_after<P: AsRef<Path>>(dir: P, id: usize) -> Result<(), DeployStatsError> {
        let dir = dir.as_ref();
        let mut find_stat = FindDeployStat::new(dir).await.map_err(|e| DeployStatsError::DirReadError(dir.to_path_buf(), e))?;
        if let Some((path, found_id)) = find_stat.next().await.map_err(|e| DeployStatsError::DirReadError(dir.to_path_buf(), e))? {
            if found_id > id {
                tokio::fs::remove_file(&path).await.map_err(|e| DeployStatsError::RemoveFileError(path, e))?;
            } 
        }
        Ok(())
    }

    async fn load<P: AsRef<Path>>(dir: P, id: usize) -> Result<Self, DeployStatsError> {
        Ok(Self {
            stats: DeployStat::load_id_range(dir, 1..=id).await?,
        })
    }

    pub async fn load_current_by_root<P: AsRef<Path>>(root: P) -> Result<Option<(Self, usize)>, DeployStatsError> {
        let dir = Self::deploy_stat_dir(root);
        Self::load_current(dir).await
    }

    pub async fn load_current<P: AsRef<Path>>(dir: P) -> Result<Option<(Self, usize)>, DeployStatsError> {
        let current_stat = Self::current_stat_path(dir.as_ref());
        let Some(current_id) = Self::read_current_id(&current_stat).await? else {
            return Ok(None);
        };
        let stats = Self::load(dir.as_ref(), current_id).await?;
        Ok(Some((stats, current_id)))
    }

    pub fn deployed_pkgs(&self) -> HashSet<&str> {
        let mut pkgs = HashSet::new();
        for (id, stat) in self.stats.iter().enumerate() {
            if stat.mode.is_some() {
                pkgs.extend(stat.pkgs.iter().map(|s| s.as_str()));
            } else {
                for pkg in stat.pkgs.iter() {
                    if !pkgs.remove(pkg.as_str()) {
                        error!("undeployed package '{}' in {}.stat not found in previous deploy stats", id, pkg);
                    }
                }
            }
        }
        pkgs
    }
}

impl FindDeployStat {
    async fn new<P: AsRef<Path>>(dir: P) -> Result<Self, tokio::io::Error> {
        Ok(Self {
            read_dir: tokio::fs::read_dir(dir).await?,
        })
    }

    pub async fn next(&mut self) -> Result<Option<(PathBuf, usize)>, std::io::Error> {
        while let Some(entry) = self.read_dir.next_entry().await? {
            let ty = entry.file_type().await?;
            if !ty.is_file() {
                continue;
            }

            let file_name = entry.file_name();
            match DeployStat::id_try_from_file_name(file_name.to_string_lossy().as_ref()) {
                Ok(id) => return Ok(Some((entry.path(), id))),
                Err(_) => {
                    warn!("invalid deploy stat file '{}'", entry.path().display());
                },
            }
        }
        Ok(None)
    }
}